import matplotlib.pyplot as plt
import numpy as np
import math
from scipy import signal
import torch.nn as nn
import torch

fs = 200.0  # 采样频率 (Hz)
# t = np.linspace(0, 0.5, 100, endpoint=False)  # 1秒钟的时间向量(1s 200个)
# x = np.sin(2 * np.pi * 10 * t) +  np.sin(2 * np.pi * 50 * t)  # 合成信号
d1 = np.loadtxt('D:/multimodal_health/dataset/Dataset EMG Fatigue/Data as txt Files/sub1.txt')
data = d1[10:24010,0]

# 高通滤波器
b2, a2 = signal.butter(8, 0.15, 'highpass')
# 截止频率为15Hz，采样频率为200Hz，Wn=2*15/200 = 0.15
filtedData2 = signal.filtfilt(b2, a2, data)

f,t,Zxx = signal.stft(filtedData2, fs,nperseg=100,noverlap=50)

plt.pcolormesh(t, f, np.abs(Zxx), shading='gouraud')
plt.colorbar(label='Magnitude')
plt.title('Gaussian Short-Time Fourier Transform (STFT)')
plt.xlabel('Time [s]')
plt.ylabel('Frequency [Hz]')
plt.show()

print(Zxx.shape)

